import gradio as gr
from colorthief import ColorThief
from PIL import Image
import io

# 颜色转换工具
def rgb_to_hex(rgb):
    return "#{:02x}{:02x}{:02x}".format(rgb[0], rgb[1], rgb[2])

# 主要功能：提取调色板
def extract_palette(image):
    # 将图像转换为内存字节流
    image_bytes = io.BytesIO()
    image.save(image_bytes, format="PNG")
    image_bytes.seek(0)
    
    # 使用 ColorThief 获取主色调调色板
    color_thief = ColorThief(image_bytes)
    palette = color_thief.get_palette(color_count=5)
    
    # 转换为十六进制表示
    hex_palette = [rgb_to_hex(color) for color in palette]
    
    # 构造输出
    result = []
    for idx, color in enumerate(palette):
        result.append(f"Color {idx + 1}: {color} | {hex_palette[idx]}")
    
    return palette, "\n".join(result)

# Gradio 界面构建
def main():
    def process(image):
        palette, details = extract_palette(image)
        # 绘制颜色块
        colors = [
            Image.new("RGB", (100, 100), color=tuple(palette[i])) for i in range(len(palette))
        ]
        color_blocks = Image.new("RGB", (500, 100))
        for i, block in enumerate(colors):
            color_blocks.paste(block, (i * 100, 0))
        return color_blocks, details

    with gr.Blocks() as app:
        gr.Markdown("## 主色调调色板分析工具")
        gr.Markdown("上传图片，获取其主色调调色板及对应的十六进制颜色代码。")

        with gr.Row():
            image_input = gr.Image(type="pil", label="上传图片")
            output_image = gr.Image(label="主色调调色板")
        
        output_text = gr.Textbox(label="颜色详情")
        submit_button = gr.Button("分析")

        submit_button.click(
            process, inputs=[image_input], outputs=[output_image, output_text]
        )

    app.launch()

if __name__ == "__main__":
    main()
